{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "# Copyright (c) Facebook, Inc. and its affiliates.\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#    http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# KubeFlow Pipelines :  HPO with AX - Pytorch Cifar10 Image classification\n",
    "\n",
    "In this example, we train a Pytorch Lightning model to using image classification cifar10 dataset. A parent run will be created during the training process,which would dump the baseline model and relevant parameters,metrics and model along with its summary,subsequently followed by a set of nested child runs, which will dump the trial results. The best parameters would be dumped into the parent run once the experiments are completed.\n",
    "\n",
    "This notebook shows PyTorch CIFAR10 end-to-end classification example using Kubeflow Pipelines. \n",
    "\n",
    "An example notebook that demonstrates how to:\n",
    "\n",
    "* Get different tasks needed for the pipeline\n",
    "* Create a Kubeflow pipeline\n",
    "* Include Pytorch KFP components to preprocess, train, visualize and deploy the model in the pipeline\n",
    "* Submit a job for execution\n",
    "* Query(prediction and explain) the final deployed model\n"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## import the necessary packages"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "source": [
    "! pip uninstall -y kfp\n",
    "! pip install --no-cache-dir kfp ax-platform"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "Found existing installation: kfp 1.0.4\n",
      "Uninstalling kfp-1.0.4:\n",
      "  Successfully uninstalled kfp-1.0.4\n",
      "Collecting kfp\n",
      "  Downloading kfp-1.6.4.tar.gz (225 kB)\n",
      "\u001b[K     |████████████████████████████████| 225 kB 7.3 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting ax-platform\n",
      "  Downloading ax_platform-0.2.0-py3-none-any.whl (815 kB)\n",
      "\u001b[K     |████████████████████████████████| 815 kB 8.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting typeguard\n",
      "  Downloading typeguard-2.12.1-py3-none-any.whl (17 kB)\n",
      "Collecting plotly\n",
      "  Downloading plotly-5.1.0-py2.py3-none-any.whl (20.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 20.6 MB 32.1 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: scipy in /opt/conda/lib/python3.8/site-packages (from ax-platform) (1.6.2)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.8/site-packages (from ax-platform) (2.11.3)\n",
      "Requirement already satisfied: scikit-learn in /opt/conda/lib/python3.8/site-packages (from ax-platform) (0.24.1)\n",
      "Requirement already satisfied: pandas in /opt/conda/lib/python3.8/site-packages (from ax-platform) (1.2.3)\n",
      "Collecting botorch==0.5.0\n",
      "  Downloading botorch-0.5.0-py3-none-any.whl (475 kB)\n",
      "\u001b[K     |████████████████████████████████| 475 kB 81.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: torch>=1.8.1 in /home/jovyan/.local/lib/python3.8/site-packages (from botorch==0.5.0->ax-platform) (1.9.0)\n",
      "Collecting gpytorch>=1.5\n",
      "  Downloading gpytorch-1.5.0-py2.py3-none-any.whl (498 kB)\n",
      "\u001b[K     |████████████████████████████████| 498 kB 84.4 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: typing-extensions in /opt/conda/lib/python3.8/site-packages (from torch>=1.8.1->botorch==0.5.0->ax-platform) (3.7.4.3)\n",
      "Collecting absl-py<=0.11,>=0.9\n",
      "  Downloading absl_py-0.11.0-py3-none-any.whl (127 kB)\n",
      "\u001b[K     |████████████████████████████████| 127 kB 90.0 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: PyYAML<6,>=5.3 in /opt/conda/lib/python3.8/site-packages (from kfp) (5.4.1)\n",
      "Requirement already satisfied: google-cloud-storage<2,>=1.20.0 in /opt/conda/lib/python3.8/site-packages (from kfp) (1.37.1)\n",
      "Requirement already satisfied: kubernetes<13,>=8.0.0 in /opt/conda/lib/python3.8/site-packages (from kfp) (10.0.1)\n",
      "Collecting google-api-python-client<2,>=1.7.8\n",
      "  Downloading google_api_python_client-1.12.8-py2.py3-none-any.whl (61 kB)\n",
      "\u001b[K     |████████████████████████████████| 61 kB 1.5 MB/s  eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: google-auth<2,>=1.6.1 in /opt/conda/lib/python3.8/site-packages (from kfp) (1.28.1)\n",
      "Requirement already satisfied: requests-toolbelt<1,>=0.8.0 in /opt/conda/lib/python3.8/site-packages (from kfp) (0.9.1)\n",
      "Requirement already satisfied: cloudpickle<2,>=1.3.0 in /opt/conda/lib/python3.8/site-packages (from kfp) (1.6.0)\n",
      "Collecting kfp-server-api<2.0.0,>=1.1.2\n",
      "  Downloading kfp-server-api-1.6.0.tar.gz (52 kB)\n",
      "\u001b[K     |████████████████████████████████| 52 kB 48.6 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: jsonschema<4,>=3.0.1 in /opt/conda/lib/python3.8/site-packages (from kfp) (3.2.0)\n",
      "Requirement already satisfied: tabulate<1,>=0.8.6 in /opt/conda/lib/python3.8/site-packages (from kfp) (0.8.9)\n",
      "Requirement already satisfied: click<8,>=7.1.1 in /opt/conda/lib/python3.8/site-packages (from kfp) (7.1.2)\n",
      "Requirement already satisfied: Deprecated<2,>=1.2.7 in /opt/conda/lib/python3.8/site-packages (from kfp) (1.2.12)\n",
      "Requirement already satisfied: strip-hints<1,>=0.1.8 in /opt/conda/lib/python3.8/site-packages (from kfp) (0.1.9)\n",
      "Collecting docstring-parser<1,>=0.7.3\n",
      "  Downloading docstring_parser-0.8.1.tar.gz (14 kB)\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h    Preparing wheel metadata ... \u001b[?25ldone\n",
      "\u001b[?25hCollecting kfp-pipeline-spec<0.2.0,>=0.1.8\n",
      "  Downloading kfp_pipeline_spec-0.1.8-py3-none-any.whl (27 kB)\n",
      "Collecting fire<1,>=0.3.1\n",
      "  Downloading fire-0.4.0.tar.gz (87 kB)\n",
      "\u001b[K     |████████████████████████████████| 87 kB 77.6 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: protobuf<4,>=3.13.0 in /opt/conda/lib/python3.8/site-packages (from kfp) (3.15.7)\n",
      "Requirement already satisfied: six in /opt/conda/lib/python3.8/site-packages (from absl-py<=0.11,>=0.9->kfp) (1.15.0)\n",
      "Requirement already satisfied: wrapt<2,>=1.10 in /opt/conda/lib/python3.8/site-packages (from Deprecated<2,>=1.2.7->kfp) (1.12.1)\n",
      "Collecting termcolor\n",
      "  Downloading termcolor-1.1.0.tar.gz (3.9 kB)\n",
      "Collecting httplib2<1dev,>=0.15.0\n",
      "  Downloading httplib2-0.19.1-py3-none-any.whl (95 kB)\n",
      "\u001b[K     |████████████████████████████████| 95 kB 72.7 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: google-api-core<2dev,>=1.21.0 in /opt/conda/lib/python3.8/site-packages (from google-api-python-client<2,>=1.7.8->kfp) (1.26.3)\n",
      "Collecting google-auth-httplib2>=0.0.3\n",
      "  Downloading google_auth_httplib2-0.1.0-py2.py3-none-any.whl (9.3 kB)\n",
      "Collecting uritemplate<4dev,>=3.0.0\n",
      "  Downloading uritemplate-3.0.1-py2.py3-none-any.whl (15 kB)\n",
      "Requirement already satisfied: googleapis-common-protos<2.0dev,>=1.6.0 in /opt/conda/lib/python3.8/site-packages (from google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (1.53.0)\n",
      "Requirement already satisfied: setuptools>=40.3.0 in /opt/conda/lib/python3.8/site-packages (from google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (49.6.0.post20210108)\n",
      "Requirement already satisfied: requests<3.0.0dev,>=2.18.0 in /opt/conda/lib/python3.8/site-packages (from google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (2.25.1)\n",
      "Requirement already satisfied: pytz in /opt/conda/lib/python3.8/site-packages (from google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (2021.1)\n",
      "Requirement already satisfied: packaging>=14.3 in /opt/conda/lib/python3.8/site-packages (from google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (20.9)\n",
      "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /opt/conda/lib/python3.8/site-packages (from google-auth<2,>=1.6.1->kfp) (4.2.1)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /opt/conda/lib/python3.8/site-packages (from google-auth<2,>=1.6.1->kfp) (0.2.8)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /opt/conda/lib/python3.8/site-packages (from google-auth<2,>=1.6.1->kfp) (4.7.2)\n",
      "Requirement already satisfied: google-resumable-media<2.0dev,>=1.2.0 in /opt/conda/lib/python3.8/site-packages (from google-cloud-storage<2,>=1.20.0->kfp) (1.2.0)\n",
      "Requirement already satisfied: google-cloud-core<2.0dev,>=1.4.1 in /opt/conda/lib/python3.8/site-packages (from google-cloud-storage<2,>=1.20.0->kfp) (1.6.0)\n",
      "Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /opt/conda/lib/python3.8/site-packages (from google-resumable-media<2.0dev,>=1.2.0->google-cloud-storage<2,>=1.20.0->kfp) (1.1.2)\n",
      "Requirement already satisfied: cffi>=1.0.0 in /opt/conda/lib/python3.8/site-packages (from google-crc32c<2.0dev,>=1.0->google-resumable-media<2.0dev,>=1.2.0->google-cloud-storage<2,>=1.20.0->kfp) (1.14.5)\n",
      "Requirement already satisfied: pycparser in /opt/conda/lib/python3.8/site-packages (from cffi>=1.0.0->google-crc32c<2.0dev,>=1.0->google-resumable-media<2.0dev,>=1.2.0->google-cloud-storage<2,>=1.20.0->kfp) (2.20)\n",
      "Requirement already satisfied: pyparsing<3,>=2.4.2 in /opt/conda/lib/python3.8/site-packages (from httplib2<1dev,>=0.15.0->google-api-python-client<2,>=1.7.8->kfp) (2.4.7)\n",
      "Requirement already satisfied: attrs>=17.4.0 in /opt/conda/lib/python3.8/site-packages (from jsonschema<4,>=3.0.1->kfp) (20.3.0)\n",
      "Requirement already satisfied: pyrsistent>=0.14.0 in /opt/conda/lib/python3.8/site-packages (from jsonschema<4,>=3.0.1->kfp) (0.17.3)\n",
      "Requirement already satisfied: urllib3>=1.15 in /opt/conda/lib/python3.8/site-packages (from kfp-server-api<2.0.0,>=1.1.2->kfp) (1.26.4)\n",
      "Requirement already satisfied: certifi in /opt/conda/lib/python3.8/site-packages (from kfp-server-api<2.0.0,>=1.1.2->kfp) (2020.12.5)\n",
      "Requirement already satisfied: python-dateutil in /opt/conda/lib/python3.8/site-packages (from kfp-server-api<2.0.0,>=1.1.2->kfp) (2.8.1)\n",
      "Requirement already satisfied: requests-oauthlib in /opt/conda/lib/python3.8/site-packages (from kubernetes<13,>=8.0.0->kfp) (1.3.0)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /opt/conda/lib/python3.8/site-packages (from kubernetes<13,>=8.0.0->kfp) (0.58.0)\n",
      "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /opt/conda/lib/python3.8/site-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.1->kfp) (0.4.8)\n",
      "Requirement already satisfied: idna<3,>=2.5 in /opt/conda/lib/python3.8/site-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (2.10)\n",
      "Requirement already satisfied: chardet<5,>=3.0.2 in /opt/conda/lib/python3.8/site-packages (from requests<3.0.0dev,>=2.18.0->google-api-core<2dev,>=1.21.0->google-api-python-client<2,>=1.7.8->kfp) (4.0.0)\n",
      "Requirement already satisfied: wheel in /opt/conda/lib/python3.8/site-packages (from strip-hints<1,>=0.1.8->kfp) (0.36.2)\n",
      "Requirement already satisfied: MarkupSafe>=0.23 in /opt/conda/lib/python3.8/site-packages (from jinja2->ax-platform) (1.1.1)\n",
      "Requirement already satisfied: numpy>=1.16.5 in /opt/conda/lib/python3.8/site-packages (from pandas->ax-platform) (1.20.2)\n",
      "Collecting tenacity>=6.2.0\n",
      "  Downloading tenacity-8.0.1-py3-none-any.whl (24 kB)\n",
      "Requirement already satisfied: oauthlib>=3.0.0 in /opt/conda/lib/python3.8/site-packages (from requests-oauthlib->kubernetes<13,>=8.0.0->kfp) (3.1.0)\n",
      "Requirement already satisfied: joblib>=0.11 in /opt/conda/lib/python3.8/site-packages (from scikit-learn->ax-platform) (1.0.1)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /opt/conda/lib/python3.8/site-packages (from scikit-learn->ax-platform) (2.1.0)\n",
      "Building wheels for collected packages: kfp, docstring-parser, fire, kfp-server-api, termcolor\n",
      "  Building wheel for kfp (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for kfp: filename=kfp-1.6.4-py3-none-any.whl size=307978 sha256=707b241c3eccfc69aef31950614cc45ff55a0d079d8321175419eacd8f9a8c2c\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-hs6eb6sp/wheels/94/27/d3/4e61074deae41e0ef5f166343fc7540eb85536c3d49f4cd25f\n",
      "  Building wheel for docstring-parser (PEP 517) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for docstring-parser: filename=docstring_parser-0.8.1-py3-none-any.whl size=19678 sha256=41b1b8ca7cfe15315cdfe91457bb99c0cd08698f3882954995eba509969c3bce\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-hs6eb6sp/wheels/7a/06/c4/9b2f0146899e8d1e7748897e42ad5412b6a025513f89cc4a0f\n",
      "  Building wheel for fire (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for fire: filename=fire-0.4.0-py2.py3-none-any.whl size=115928 sha256=0bc93566fe7b907931734eb71451051e44437d1f14f3b0cacf56c835e0cad3fc\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-hs6eb6sp/wheels/1f/10/06/2a990ee4d73a8479fe2922445e8a876d38cfbfed052284c6a1\n",
      "  Building wheel for kfp-server-api (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for kfp-server-api: filename=kfp_server_api-1.6.0-py3-none-any.whl size=92524 sha256=6395d6cd9d9e57f03258b9c515ecadff2e5b105609f1cb96d618bd08ba6213e2\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-hs6eb6sp/wheels/73/72/bd/804c3bcea41a99dc328a623d003ba9d7d243f2a59a767ef28c\n",
      "  Building wheel for termcolor (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for termcolor: filename=termcolor-1.1.0-py3-none-any.whl size=4829 sha256=64bc3882cc685e2dfbbec60ec1ae3b38f907e52100c137c3494a0357673a7d42\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-hs6eb6sp/wheels/a0/16/9c/5473df82468f958445479c59e784896fa24f4a5fc024b0f501\n",
      "Successfully built kfp docstring-parser fire kfp-server-api termcolor\n",
      "Installing collected packages: httplib2, uritemplate, termcolor, tenacity, gpytorch, google-auth-httplib2, typeguard, plotly, kfp-server-api, kfp-pipeline-spec, google-api-python-client, fire, docstring-parser, botorch, absl-py, kfp, ax-platform\n",
      "  Attempting uninstall: kfp-server-api\n",
      "    Found existing installation: kfp-server-api 1.0.4\n",
      "    Uninstalling kfp-server-api-1.0.4:\n",
      "      Successfully uninstalled kfp-server-api-1.0.4\n",
      "Successfully installed absl-py-0.11.0 ax-platform-0.2.0 botorch-0.5.0 docstring-parser-0.8.1 fire-0.4.0 google-api-python-client-1.12.8 google-auth-httplib2-0.1.0 gpytorch-1.5.0 httplib2-0.19.1 kfp-1.6.4 kfp-pipeline-spec-0.1.8 kfp-server-api-1.6.0 plotly-5.1.0 tenacity-8.0.1 termcolor-1.1.0 typeguard-2.12.1 uritemplate-3.0.1\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "import kfp\n",
    "import json\n",
    "import os\n",
    "from kfp.onprem import use_k8s_secret\n",
    "from kfp import components\n",
    "from kfp.components import load_component_from_file, load_component_from_url, func_to_container_op, InputPath\n",
    "from kfp import dsl\n",
    "from kfp import compiler\n",
    "\n",
    "import numpy as np\n",
    "import logging\n",
    "\n",
    "from ax.service.ax_client import AxClient\n",
    "import json\n",
    "\n",
    "kfp.__version__"
   ],
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "'1.6.4'"
      ]
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Enter your gateway and the auth token\n",
    "[Use this extension on chrome to get token]( https://chrome.google.com/webstore/detail/editthiscookie/fngmhnnpilhplaeedifhccceomclgfbg?hl=en)\n",
    "\n",
    "![image.png](./image.png)"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Update values for the ingress gateway and auth session"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "INGRESS_GATEWAY='http://istio-ingressgateway.istio-system.svc.cluster.local'\n",
    "AUTH=\"<auth-token>\"  \n",
    "NAMESPACE=\"kubeflow-user-example-com\"\n",
    "COOKIE=\"authservice_session=\"+AUTH\n",
    "EXPERIMENT=\"Default\"\n",
    "dist_volume = 'dist-vol'\n",
    "volume_mount_path =\"/model\"\n",
    "results_path = volume_mount_path+\"/results.json\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Set  the Log bucket and Tensorboard Image"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "MINIO_ENDPOINT=\"http://minio-service.kubeflow:9000\"\n",
    "LOG_BUCKET=\"mlpipeline\"\n",
    "TENSORBOARD_IMAGE=\"public.ecr.aws/pytorch-samples/tboard:latest\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Set the client and create the experiment"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "client = kfp.Client(host=INGRESS_GATEWAY+\"/pipeline\", cookies=COOKIE)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "client.create_experiment(EXPERIMENT)\n",
    "experiments = client.list_experiments(namespace=NAMESPACE)\n",
    "my_experiment = experiments.experiments[0]\n",
    "my_experiment"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "<a href=\"http://istio-ingressgateway.istio-system.svc.cluster.local/pipeline/#/experiments/details/ba9b7266-2b1c-4729-afcd-be808c25c5af\" target=\"_blank\" >Experiment details</a>."
      ]
     },
     "metadata": {}
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "{'created_at': datetime.datetime(2021, 6, 21, 13, 13, 6, tzinfo=tzlocal()),\n",
       " 'description': None,\n",
       " 'id': 'ba9b7266-2b1c-4729-afcd-be808c25c5af',\n",
       " 'name': 'Default',\n",
       " 'resource_references': [{'key': {'id': 'kubeflow-user-example-com',\n",
       "                                  'type': 'NAMESPACE'},\n",
       "                          'name': None,\n",
       "                          'relationship': 'OWNER'}],\n",
       " 'storage_state': 'STORAGESTATE_AVAILABLE'}"
      ]
     },
     "metadata": {},
     "execution_count": 6
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Set  the Inference parameters"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "DEPLOY_NAME=\"torchserve\"\n",
    "MODEL_NAME=\"cifar10\"\n",
    "ISVC_NAME=DEPLOY_NAME+\".\"+NAMESPACE+\".\"+\"example.com\"\n",
    "INPUT_REQUEST=\"https://raw.githubusercontent.com/kubeflow/pipelines/master/samples/contrib/pytorch-samples/cifar10/input.json\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Load the the components yaml files for setting up the components"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "source": [
    "prepare_tensorboard_op = load_component_from_file(\"./common/tensorboard/component.yaml\")\n",
    "\n",
    "generate_trails_op = components.load_component_from_file(\n",
    "    \"./common/ax_generate_trials/component.yaml\"\n",
    ")\n",
    "\n",
    "complete_trails_op = components.load_component_from_file(\n",
    "    \"./common/ax_complete_trials/component.yaml\"\n",
    ")\n",
    "\n",
    "get_keys_op = components.load_component_from_file(\n",
    "    \"../../../components/json/Get_keys/component.yaml\"\n",
    ")\n",
    "\n",
    "get_element_op = components.load_component_from_file(\n",
    "    \"../../../components/json/Get_element_by_key/component.yaml\"\n",
    ")\n",
    "prep_op = components.load_component_from_file(\n",
    "    \"./cifar10/yaml/pre_process/component.yaml\"\n",
    ")\n",
    "\n",
    "# Uncomment hpo inputs in component yaml\n",
    "train_op = components.load_component_from_file(\n",
    "    \"./cifar10/yaml/train/component.yaml\"\n",
    ")\n",
    "\n",
    "deploy_op = load_component_from_file(\"./common/deploy/component.yaml\")\n",
    "\n",
    "pred_op = load_component_from_file(\"./common/prediction/component.yaml\")\n",
    "\n",
    "minio_op = components.load_component_from_file(\n",
    "    \"./common/minio/component.yaml\"\n",
    ")\n",
    "\n",
    "kubernetes_create_pvc_op = load_component_from_file(\"../../../components/kubernetes/Create_PersistentVolumeClaim/component.yaml\")"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "from kubernetes.client.models import V1Volume, V1PersistentVolumeClaimVolumeSource\n",
    "def create_dist_pipeline():\n",
    "    kubernetes_create_pvc_op(name=dist_volume, storage_size= \"20Gi\")\n",
    "\n",
    "create_volume_run = client.create_run_from_pipeline_func(create_dist_pipeline, arguments={})\n",
    "create_volume_run.wait_for_run_completion()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "source": [
    "parameters = [\n",
    "                {\"name\": \"lr\", \"type\": \"range\", \"bounds\": [1e-4, 0.2], \"log_scale\": True},\n",
    "                {\"name\": \"weight_decay\", \"type\": \"range\", \"bounds\": [1e-4, 1e-2]},\n",
    "                {\"name\": \"eps\", \"type\": \"range\", \"bounds\": [1e-8, 1e-2]},\n",
    "             ]"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Define the pipeline"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "source": [
    "@dsl.pipeline(\n",
    "    name=\"AX Hpo\", description=\"Estimating best parameters using AX\"\n",
    ")\n",
    "def pytorch_ax_hpo( # pylint: disable=too-many-arguments\n",
    "    minio_endpoint=MINIO_ENDPOINT,\n",
    "    log_bucket=LOG_BUCKET,\n",
    "    log_dir=f\"tensorboard/logs/{dsl.RUN_ID_PLACEHOLDER}\",\n",
    "    mar_path=f\"mar/{dsl.RUN_ID_PLACEHOLDER}/model-store\",\n",
    "    config_prop_path=f\"mar/{dsl.RUN_ID_PLACEHOLDER}/config\",\n",
    "    model_uri=f\"s3://mlpipeline/mar/{dsl.RUN_ID_PLACEHOLDER}\",\n",
    "    best_params=f\"hpo/{dsl.RUN_ID_PLACEHOLDER}\",\n",
    "    tf_image=TENSORBOARD_IMAGE,\n",
    "    deploy=DEPLOY_NAME,\n",
    "    isvc_name=ISVC_NAME,\n",
    "    model=MODEL_NAME,\n",
    "    namespace=NAMESPACE,\n",
    "    confusion_matrix_log_dir=f\"confusion_matrix/{dsl.RUN_ID_PLACEHOLDER}/\",\n",
    "    checkpoint_dir=\"checkpoint_dir/cifar10\",\n",
    "    input_req=INPUT_REQUEST,\n",
    "    cookie=COOKIE,\n",
    "    total_trials=2,\n",
    "    ingress_gateway=INGRESS_GATEWAY,\n",
    "):\n",
    "   \n",
    "    \"\"\"This method defines the pipeline tasks and operations\"\"\"\n",
    "    pod_template_spec = json.dumps({\n",
    "        \"spec\": {\n",
    "            \"containers\": [{\n",
    "                \"env\": [\n",
    "                    {\n",
    "                        \"name\": \"AWS_ACCESS_KEY_ID\",\n",
    "                        \"valueFrom\": {\n",
    "                            \"secretKeyRef\": {\n",
    "                                \"name\": \"mlpipeline-minio-artifact\",\n",
    "                                \"key\": \"accesskey\",\n",
    "                            }\n",
    "                        },\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"AWS_SECRET_ACCESS_KEY\",\n",
    "                        \"valueFrom\": {\n",
    "                            \"secretKeyRef\": {\n",
    "                                \"name\": \"mlpipeline-minio-artifact\",\n",
    "                                \"key\": \"secretkey\",\n",
    "                            }\n",
    "                        },\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"AWS_REGION\",\n",
    "                        \"value\": \"minio\"\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"S3_ENDPOINT\",\n",
    "                        \"value\": f\"{minio_endpoint}\",\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"S3_USE_HTTPS\",\n",
    "                        \"value\": \"0\"\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"S3_VERIFY_SSL\",\n",
    "                        \"value\": \"0\"\n",
    "                    },\n",
    "                ]\n",
    "            }]\n",
    "        }\n",
    "    })\n",
    "\n",
    "    prepare_tb_task = prepare_tensorboard_op(\n",
    "        log_dir_uri=f\"s3://{log_bucket}/{log_dir}\",\n",
    "        image=tf_image,\n",
    "        pod_template_spec=pod_template_spec,\n",
    "    ).set_display_name(\"Visualization\")\n",
    "\n",
    "    prep_task = (\n",
    "        prep_op().after(prepare_tb_task).set_display_name(\"Preprocess & Transform\")\n",
    "    )\n",
    "\n",
    "    gen_trials_task = generate_trails_op(total_trials, parameters, 'test-accuracy').after(prep_task).set_display_name(\"AX Generate Trials\")\n",
    "    \n",
    "    get_keys_task = get_keys_op(gen_trials_task.outputs[\"trial_parameters\"]).after(gen_trials_task).set_display_name(\"Get Keys of Trials\")\n",
    "    \n",
    "    confusion_matrix_url = f\"minio://{log_bucket}/{confusion_matrix_log_dir}\"\n",
    "    script_args = f\"model_name=resnet.pth,\" \\\n",
    "                  f\"confusion_matrix_url={confusion_matrix_url}\"\n",
    "    ptl_args = f\"max_epochs=1, profiler=pytorch\"\n",
    "\n",
    "    with dsl.ParallelFor(get_keys_task.outputs[\"keys\"]) as item:\n",
    "        get_element_task = get_element_op(gen_trials_task.outputs[\"trial_parameters\"], item).after(get_keys_task).set_display_name(\"Get Element from key\")\n",
    "        train_task = (\n",
    "            train_op(\n",
    "                trial_id=item,\n",
    "                input_data=prep_task.outputs[\"output_data\"],\n",
    "                cifar_script_args=script_args,\n",
    "                model_parameters=get_element_task.outputs[\"output\"],\n",
    "                ptl_arguments=ptl_args,\n",
    "                results=results_path\n",
    "            ).add_pvolumes({volume_mount_path: dsl.PipelineVolume(pvc=dist_volume)}).after(get_element_task).set_display_name(\"Training\")\n",
    "#             For GPU uncomment below line and set GPU limit and node selector\n",
    "#             ).set_gpu_limit(1).add_node_selector_constraint('cloud.google.com/gke-accelerator','nvidia-tesla-p4')\n",
    "        )\n",
    "        \n",
    "    complete_trials_task = complete_trails_op(gen_trials_task.outputs[\"client\"], results_path).add_pvolumes({volume_mount_path: dsl.PipelineVolume(pvc=dist_volume)}).after(train_task).set_display_name(\"AX Complete Trials\")\n",
    "\n",
    "    dsl.get_pipeline_conf().add_op_transformer(\n",
    "        use_k8s_secret(\n",
    "            secret_name=\"mlpipeline-minio-artifact\",\n",
    "            k8s_secret_key_to_env={\n",
    "                \"secretkey\": \"MINIO_SECRET_KEY\",\n",
    "                \"accesskey\": \"MINIO_ACCESS_KEY\",\n",
    "            },\n",
    "        )\n",
    "    )\n"
   ],
   "outputs": [],
   "metadata": {
    "tags": []
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Compile  the pipeline"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "source": [
    "compiler.Compiler().compile(pytorch_ax_hpo, 'pytorch.tar.gz', type_check=True)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Execute the pipeline"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "source": [
    "run = client.run_pipeline(my_experiment.id, 'pytorch_ax_hpo', 'pytorch.tar.gz')"
   ],
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "<a href=\"http://istio-ingressgateway.istio-system.svc.cluster.local/pipeline/#/runs/details/f848da24-d24d-4100-a0de-dd9a089f0435\" target=\"_blank\" >Run details</a>."
      ]
     },
     "metadata": {}
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}